"""
mathmodeling_SAK.Analysis.trend_testing
=======================================

对数据序列进行趋势性检验，查看数据是否具有增长或者消减的趋势。这些结果对
于一些对数据趋势性敏感的分析方法，如 ARIMA 模型，至关重要。
"""

# 导入必要的库
import pandas as pd
import numpy as np

from scipy.stats import linregress
import pymannkendall as mk

from statsmodels.tsa.stattools import adfuller

def autoADF(data_series):
    r"""
    执行增强单位根检验，即 Dickey Fuller（ADF）测试，以检验数据的平
    稳性。检验结果会以 `pandas.DataFrame` 的形式返回。
    
    .. note::
    
        ADF测试假设时间序列是均匀的。确保对数据进行相应的预处理。

    输入参数
    ----------
    data_series : pd.Series
        有待测试的时间序列数据

    返回值
    ----------
    pd.DataFrame
        包含 ADF 测试结果的 DataFrame，包含以下行：
        
        - "Test Statistic Value": 测试统计值。
        - "p-value": 测试的 p 值。通常用于与一个预先给出的置信度指标
          :math:`\alpha` 进行大小的比较。如果 p-values 小于这个数
          值，我们就认为数据序列趋势是平稳的。否则不平稳。
        - "Lags Used": 测试中使用的滞后次数。
        - "Number of Observations Used": 测试中使用的观察数量。
        - "Critical Value(1%)": 在1%显著性水平上的临界值。
        - "Critical Value(5%)": 在5%显著性水平上的临界值。
        - "Critical Value(10%)": 在10%显著性水平上的临界值。
        
    .. note::
    
        返回的 `pandas.DataFrame` 是纵向的，指标参数的名称陈列在表
        格左侧的一列中。
    
    """
    t = adfuller(data_series) # ADF检验

    result = pd.DataFrame(  
            index = [
                "Test Statistic Value", 
                "p-value", 
                "Lags Used", 
                "Number of Observations Used",
                "Critical Value(1%)",
                "Critical Value(5%)",
                "Critical Value(10%)" 
                ],
            columns=['value'] 
        )

    result.loc['value', 'Test Statistic Value'] = t[0]
    result.loc['value', 'p-value'] = t[1]
    result.loc['value', 'Lags Used'] = t[2]
    result.loc['value', 'Number of Observations Used'] = t[3]
    result.loc['value', 'Critical Value(1%)'] = t[4]['1%']
    result.loc['value', 'Critical Value(5%)'] = t[4]['5%']
    result.loc['value', 'Critical Value(10%)'] = t[4]['10%']

    return result

def n_diff(data, n=1):
    """
    对数据的某一列进行 n 阶差分，并将结果作为新列追加到原始数据后面。

    输入参数
    ----------
    data : pd.DataFrame
        包含待差分列的数据框。
    n : int, optional
        差分的阶数 (default is 1)。

    返回值
    ----------
    pd.DataFrame
        添加了差分列的数据框。
    """
    diff_col_name = f'A_diff_{n}'
    data[diff_col_name] = np.nan
    data[diff_col_name][n:] = np.diff(data['A'], n=n)
    return data

def cox_stuart_test(data):
    """
    Cox-Stuart 数据趋势检验检验
    
    直接考虑数据的变化趋势，若数据有上升趋势，那么排在后面的数据的值要比
    排在前面的数据的值显著的大，反之，若数据有下降趋势，那么排在后面的数
    据的值要比排在前面的数据的值显著的小，利用前后两个时期不同数据的差值
    正负来判断数据总的变化趋势。

    输入参数
    ----------
    data : pd.Series
        需要检验趋势的数据序列。为时间序列数据。

    返回值
    ----------
    p_value
        检验结果的 p 值
    trend, str
        趋势检验的结果
    """
    n = len(data)
    half = n // 2
    first_half = data[:half]
    second_half = data[-half:]
    signs = np.sign(second_half - first_half)
    pos_count = np.sum(signs > 0)
    neg_count = np.sum(signs < 0)
    p_value = 2 * min(pos_count, neg_count) / half
    trend = 'increasing' if (
        pos_count > neg_count
        ) else 'decreasing'
    return p_value, trend

def mann_kendall_test(data):
    """
    Mann-Kendall 数据趋势检验 
    
    Mann-kendall检验不需要样本遵循一定的分布，也不说少数异常值的干扰。
    对于统计量z，大于0时时上升趋势，小于0时是下降趋势。

    该检验的目的是统计评估我们所感兴趣的变量，随着时间的变化，是否有单调
    上升或下降的趋势。单调上升或者下降的趋势意味着该变量随着时间增加或者
    减少。但此趋势可能是，也可能不是线性的，MK test可替代参数线性回归分
    析——线性回归可检测线性拟合直线是否不为零。回归分析要求拟合回归线的残
    差是否是正太分布的。MK检验不需要这种假设，MK检验是非参数检验（不要求
    服从任何分布distribution free）,某论文表明MK检验最好被视作探索性
    分析，最适合用于识别变化显著或浮动较大的站点，并量化这些结果。

    输入参数
    ----------
    data : pd.Series
        需要检验趋势的数据序列。为时间序列数据。

    返回值
    ----------
    p_value
        检验结果的 p 值
    trend, str
        趋势检验的结果
    """
    result = mk.original_test(data)
    p_value = result.p
    trend = result.trend
    return p_value, trend

def slope_test(data):
    """
    使用简单的斜率方法对数据进行趋势性检验。

    输入参数
    ----------
    data : pd.Series
        需要检验趋势的数据序列。为时间序列数据。

    返回值
    ----------
    p_value
        检验结果的 p 值
    trend, str
        趋势检验的结果
    """
    slope, intercept, r_value, p_value, std_err = linregress(
        range(len(data)), 
        data
        )
    trend = 'increasing' if slope > 0 else 'decreasing'
    return p_value, trend

def trend_tests_df(data):
    """
    一次性使用 Cox-Stuart, Mann-Kendall, 以及简单的斜率方法三种方
    法对数据的趋势性进行检验。

    输入参数
    ----------
    data : list or np.array
        The time series data to be tested.

    返回值
    ----------
    pd.DataFrame
        A DataFrame containing the p-values and trend direction for each test.
    """
    data = np.array(data)
    
    # Cox-Stuart Test
    (
        cox_stuart_p_value, 
        cox_stuart_trend
        ) = cox_stuart_test(data)
    
    # Mann-Kendall Test
    (
        mann_kendall_p_value, 
        mann_kendall_trend
        ) = mann_kendall_test(data)
    
    # Slope Test
    (slope_p_value, slope_trend) = slope_test(data)
    
    # Compile results into a DataFrame
    results = pd.DataFrame({
        'Test': [
            'Cox-Stuart', 
            'Mann-Kendall', 
            'Slope'],
        'p-value': [
            cox_stuart_p_value, 
            mann_kendall_p_value, 
            slope_p_value],
        'Trend': [
            cox_stuart_trend, 
            mann_kendall_trend, 
            slope_trend]
    })
    
    return results

def iterative_trend_tests(
    data, 
    stop_test='Mann-Kendall', 
    max_iter_of_trend_test=5):
    r"""
    检验数据序列的趋势。如果数据序列存在趋势，就进行差分，再次检验，直至
    不再存在趋势，或者达到了 `max_iter_of_trend_test` 就停止迭代并
    退出循环，保存每一次的计算结果，最终以 DataFrame 的形式输出。

    输入参数
    ----------
    data : list or np.array
        The time series data to be tested.
    stop_test : str, optional
        停止迭代所依赖的检验方法，可选参数有 'All', 'Cox-Stuart', 
       'Mann-Kendall', 'Slope'。如果选择了 'All' 则要求所有检验都
       认为数据无趋势时，迭代才会停止。default by 'Mann-Kendall'
    max_iter_of_trend_test : int, optional
        测试的最大迭代次数，default by 5.

    返回值
    ----------
    pd.DataFrame
        包含各个差分阶数下三种数据趋势检验方法的 p-values 等指标的结
        果表。
    """
    results = []
    iteration = 0
    while iteration < max_iter_of_trend_test:
        test_results = trend_tests_df(data)
        test_results['Iteration'] = iteration
        results.append(test_results)
        
        if stop_test == 'All':
            if all(test_results['Trend'] == 'no trend'):
                break
        else:
            stop_test_result = test_results[
                test_results['Test'] == stop_test]
            if stop_test_result[
                'Trend'
                ].values[0] == 'no trend':
                break
        
        data = np.diff(data)
        iteration += 1
    
    return pd.concat(results, ignore_index=True)